package com.googlecode.yoohoo.io.socket.stream;

import java.io.StringReader;
import java.io.UnsupportedEncodingException;
import java.util.HashMap;
import java.util.Map;

import javax.security.sasl.Sasl;
import javax.security.sasl.SaslException;
import javax.security.sasl.SaslServer;
import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamConstants;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;

import org.apache.commons.codec.binary.Base64;

import com.googlecode.yoohoo.xmppcore.model.account.IAccountManager;
import com.googlecode.yoohoo.xmppcore.protocol.XmppProtocolException;
import com.googlecode.yoohoo.xmppcore.protocol.stream.error.InvalidXml;
import com.googlecode.yoohoo.xmppcore.protocol.stream.error.XmlNotWellFormed;
import com.googlecode.yoohoo.xmppcore.stream.INegotiationContext;

public class SaslNegotiant extends InitialStreamNegotiant {
	enum SaslNegotiationState {
		NONE,
		AUTH_REQUEST_RECEIVED,
		CHALLENGE_SENT,
		RESPONSE_RECEIVED,
		ABORTED,
		FAILED,
		SUCCESSFUL
	};
	
	// rfc3920(6.4)
	enum SaslErrorCondition {
		//Aborted, // need process especially
		IncorrectEncoding,
		InvalidAuthzid,
		InvalidMechanism,
		MechanismTooWeak,
		NotAuthorized,
		TemporaryAuthFailure
	};
	
	private SaslNegotiationState state;
	private static final String XMPP_SASL_AUTH_TAG_NAME = "auth";
	private static final String XMPP_SASL_NAMESPACE = "urn:ietf:params:xml:ns:xmpp-sasl";
	private static final String XMPP_SASL_RESPONSE_TAG_NAME = "response";
	private int abortRetries;
	private int failureRetries;
	private SaslServer saslServer;
	private static Base64 base64;
	private boolean authMaxRetriesHasSet;
	private int authMaxRetries;
	private String hostName;
	private IAccountManager accountManager;
	
	static {
		base64 = new Base64();
	}
	
	public SaslNegotiant(IStreamFeaturesManager featuresManager, IAccountManager accountManager,
			String hostName, int authMaxRetries) {
		super(featuresManager);
		state = SaslNegotiationState.NONE;
		abortRetries = 0;
		failureRetries = 0;
		this.hostName = hostName;
		this.authMaxRetries = authMaxRetries;
		this.accountManager = accountManager;
	}

	@Override
	protected void doNegotiate(INegotiationContext context, String message) {
		if (!authMaxRetriesHasSet) {
			if (authMaxRetries < 2)
				authMaxRetries = 2;
			
			abortRetries = authMaxRetries;
			failureRetries = authMaxRetries;
			authMaxRetriesHasSet = true;
		}
		
		try {
			saslNegotiate(context, message);
		} catch (Exception e) {
			processFailure(context, SaslErrorCondition.TemporaryAuthFailure);
		}
	}
	
	private void processFailure(INegotiationContext context, SaslErrorCondition condition) {
		String xmlTemplate = "<failure xmlns='urn:ietf:params:xml:ns:xmpp-sasl'>%s</failure>";
		String xml = null;
		
		if (condition == SaslErrorCondition.IncorrectEncoding) {
			xml = String.format(xmlTemplate, "<incorrect-encoding/>");
		} else if (condition == SaslErrorCondition.InvalidAuthzid) {
			xml = String.format(xmlTemplate, "<invalid-authzid/>");
		} else if (condition == SaslErrorCondition.InvalidMechanism) {
			xml = String.format(xmlTemplate, "<invalid-mechanism/>");
		} else if (condition == SaslErrorCondition.MechanismTooWeak) {
			xml = String.format(xmlTemplate, "<mechanism-too-weak/>");
		} else if (condition == SaslErrorCondition.NotAuthorized) {
			xml = String.format(xmlTemplate, "<not-authorized/>");
		} else if (condition == SaslErrorCondition.TemporaryAuthFailure) {
			xml = String.format(xmlTemplate, "<temporary-auth-failure/>");
		}
		
		context.write(xml);
		if (failureRetries == 0) {
			context.close();
		} else {
			failureRetries--;
		}
		state = SaslNegotiationState.FAILED;
	}

	private void saslNegotiate(INegotiationContext context, String message) {
		if (state == SaslNegotiationState.NONE ||
				state == SaslNegotiationState.ABORTED ||
					state == SaslNegotiationState.FAILED) {
			String mechanism = processAuth(context, message);
			
			if (!"DIGEST-MD5".equals(mechanism)) {
				processFailure(context, SaslErrorCondition.InvalidMechanism);
				return;
			}
			
			state = SaslNegotiationState.AUTH_REQUEST_RECEIVED;
			try {
				Map<String, Object> properties = new HashMap<String, Object>();
				properties.put(Sasl.QOP, "auth");
				saslServer = Sasl.createSaslServer("DIGEST-MD5", "xmpp", hostName, properties,
						new DefaultAuthenticationCallbackHandler((ResourceBindingNegotiant)nextNegotiant,
								accountManager));
				String challengeXml = null;
				challengeXml = createChallengeXml(new byte[0]);
				if (challengeXml != null)
					context.write(challengeXml);
				state = SaslNegotiationState.CHALLENGE_SENT;
			} catch (SaslException e) {
				processFailure(context, SaslErrorCondition.NotAuthorized);
				return;
			} catch (UnsupportedEncodingException e) {
				processFailure(context, SaslErrorCondition.IncorrectEncoding);
				return;
			}
		} else if (state == SaslNegotiationState.CHALLENGE_SENT) {
			if (message == "<abort/>") {
				processAbort(context);
				return;
			}
			
			byte[] response = null;
			try {
				response = processResponse(context, message);
			} catch (UnsupportedEncodingException e) {
				processFailure(context, SaslErrorCondition.IncorrectEncoding);
				return;
			}
			
			state = SaslNegotiationState.RESPONSE_RECEIVED;
			String challengeXml = null;
			try {
				challengeXml = createChallengeXml(response);
			} catch (SaslException e) {
				processFailure(context, SaslErrorCondition.NotAuthorized);
				return;
			} catch (UnsupportedEncodingException e) {
				processFailure(context, SaslErrorCondition.TemporaryAuthFailure);
				return;
			}
			if (challengeXml != null) {
				context.write(challengeXml);
				state = SaslNegotiationState.CHALLENGE_SENT;
			} else {
				context.write("<success xmlns='urn:ietf:params:xml:ns:xmpp-sasl'/>");
				state = SaslNegotiationState.SUCCESSFUL;
			}
		} else if (state == SaslNegotiationState.SUCCESSFUL) {
			super.doNegotiate(context, message);
		}
	}

	private void processAbort(INegotiationContext context) {
		context.write("<failure xmlns='urn:ietf:params:xml:ns:xmpp-sasl'></aborted></failure>");
		if (abortRetries == 0) {
			context.close();
		} else {
			abortRetries--;
		}
		state = SaslNegotiationState.ABORTED;
	}

	private String processAuth(INegotiationContext context, String message) {
		XMLInputFactory factory = XMLInputFactory.newInstance();
		try {
			XMLStreamReader reader = factory.createXMLStreamReader(new StringReader(message));
			int eventType = reader.next();
			
			if (eventType != XMLStreamConstants.START_ELEMENT) {
				throw new XmppProtocolException(new InvalidXml());
			}
			
			if (!XMPP_SASL_AUTH_TAG_NAME.equals(reader.getLocalName()))
				throw new XmppProtocolException(new InvalidXml());
			
			if (reader.getNamespaceCount() != 1)
				throw new XmppProtocolException(new InvalidXml());
			
			if (!XMPP_SASL_NAMESPACE.equals(reader.getNamespaceURI(0)))
				throw new XmppProtocolException(new InvalidXml());
			
			if (reader.getAttributeCount() != 1 ||
					!reader.getAttributeLocalName(0).equals("mechanism") ||
					!reader.getAttributePrefix(0).equals(""))
				throw new XmppProtocolException(new InvalidXml());
			
			String mechanism = reader.getAttributeValue(0);
			
			eventType = reader.next();
			if (eventType != XMLStreamConstants.END_ELEMENT)
				throw new XmppProtocolException(new InvalidXml());
			
			eventType = reader.next();
			if (eventType != XMLStreamConstants.END_DOCUMENT)
				throw new XmppProtocolException(new InvalidXml());
			
			return mechanism;
		} catch (XMLStreamException e) {
			throw new XmppProtocolException(new XmlNotWellFormed());
		}
	}

	private String createChallengeXml(byte[] response) throws SaslException, UnsupportedEncodingException {
		if (saslServer.isComplete())
			return null;
		
		byte[] challenge = saslServer.evaluateResponse(response);
		byte[] encodedChallenge = base64.encode(challenge);
			
		return "<challenge xmlns='urn:ietf:params:xml:ns:xmpp-sasl'>" +
					new String(encodedChallenge, "Latin1") + "</challenge>";
	}

	private byte[] processResponse(INegotiationContext context, String message) throws UnsupportedEncodingException {
		XMLInputFactory factory = XMLInputFactory.newInstance();
		try {
			XMLStreamReader reader = factory.createXMLStreamReader(new StringReader(message));
			int eventType = reader.next();
			
			if (eventType != XMLStreamConstants.START_ELEMENT) {
				throw new XmppProtocolException(new InvalidXml());
			}
			
			if (!XMPP_SASL_RESPONSE_TAG_NAME.equals(reader.getLocalName()))
				throw new XmppProtocolException(new InvalidXml());
			
			if (reader.getNamespaceCount() != 1)
				throw new XmppProtocolException(new InvalidXml());
			
			if (!XMPP_SASL_NAMESPACE.equals(reader.getNamespaceURI(0)))
				throw new XmppProtocolException(new InvalidXml());
			
			String base64Response = "";
			eventType = reader.next();
			if (eventType == XMLStreamConstants.CHARACTERS) {
				base64Response = reader.getText();
				eventType = reader.next();
			}
			
			if (eventType != XMLStreamConstants.END_ELEMENT) {
				throw new XmppProtocolException(new InvalidXml());				
			}
			
			eventType = reader.next();
			if (eventType != XMLStreamConstants.END_DOCUMENT)
				throw new XmppProtocolException(new InvalidXml());
			
			return base64.decode(base64Response.getBytes("Latin1"));
		} catch (XMLStreamException e) {
			throw new XmppProtocolException(new XmlNotWellFormed());
		}
	}

	@Override
	protected NegotiationPhase getPhase() {
		return NegotiationPhase.SASL;
	}
}
